import json
import faiss
import numpy as np
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 将一维的文本数据通过模型转成二维向量
model_id = "nlp_corom_sentence-embedding_chinese-base-medical"
pipeline_se = pipeline(Tasks.sentence_embedding, model=model_id)

def createVector():
    texts = []
    with open("1.txt", "r", encoding="utf-8") as f:
        inputs = json.loads(f.read()).get("dataList")

        for data in inputs:
            texts.append(data.get("text"))
        inputs = {
            'source_sentence': texts,
        }

        result = pipeline_se(input=inputs)
        embeddings = result.get("text_embedding")

        # 将结果转换为 NumPy 数组
        embeddings_array = np.array(embeddings).astype('float32')

        # 创建 FAISS 索引并添加向量
        index = faiss.IndexFlatL2(embeddings_array.shape[1])  # 维度是 embedding 的第二维
        print("创建索引：" + str(index))
        index.add(embeddings_array)  # 添加 NumPy 数组到索引

        faiss.write_index(index, "./faiss_medical_test.index")